from pydantic import BaseModel
import datetime
import pandas as pd
import numpy as np
import shutil
from pathlib import Path
from gen_html import gen_html
from generate_3h_kline import generate_3h_kline
import os
import warnings

warnings.filterwarnings("ignore")


class Performance(BaseModel):
    """
    用于记录策略表现
    """
    annualized_return: float = 0.0  # 年化收益率
    annualized_return_no_fee: float = 0.0  # 无手续费的年化收益率
    annualized_volatility: float = 0.0  # 年化波动率
    annualized_shape: float = 0.0  # 夏普比率
    turnover: float = 0.0  # 换手率, 是日均口径换手率
    margin: float = 0.0  # 代表平均每次收益获利水平，衡量盈利是否覆盖手续费
    fitness: float = 0.0  # 综合衡量夏普比率和Margin的指标
    excess_return: float = 0.0  # 年化超额收益
    relative_return: float = 0.0  # 年化相对收益倍数
    information_ratio: float = 0.0  # 信息比率，衡量超额收益的综合指标
    max_drawdown: float = 0.0  # 最大回撤
    calmar_ratio: float = 0.0  # 综合衡量年化收益和最大回撤的指标
    yearly_count: float = 0.0  # 平均年交易次数
    winning_rate: float = 0.0  # 胜率
    profit_loss_ratio: float = 0.0  # 盈亏比
    kelly_fraction: float = 0.0  # 用凯利公式计算的合理配置仓位

    @staticmethod
    def get_all_columns():
        # 返回所有指标
        return ["annualized_shape", "annualized_return", "annualized_return_no_fee", "excess_return", "calmar_ratio",
                "yearly_count", "margin", "winning_rate", "profit_loss_ratio",
                "kelly_fraction", "annualized_volatility",
                "turnover", "fitness", "relative_return",
                "max_drawdown", "information_ratio"]

    @staticmethod
    def get_columns_for_line_chart():
        # 返回画净值曲线top2的指标
        return ["annualized_shape", "margin", "information_ratio", "calmar_ratio", "kelly_fraction"]

    @staticmethod
    def get_columns_for_annual_display():
        # 返回分年度展示时的指标
        return ["excess_return", "max_drawdown", "margin"]


class RelativeAmplitudeStrategy(object):
    """
    用std计算上下轨, 突破上轨买入, 回落中轨卖出
    """

    def __init__(self, frequency):
        self.frequency = frequency  # 回测频率
        self.begin_date_list = ["20170101", "20200101", "20230101"]  # 起始日期列表
        self.data_begin_date = datetime.date(2015, 1, 1)  # 数据开始日期 数据开始日期小于回测开始日期 预留一些数据计算指标
        self.backtest_begin_date = datetime.date(2017, 1, 1)  # 回测开始日期
        self.backtest_end_date = datetime.date(2025, 8, 31)  # 回测结束日期
        self.get_R_Th_list()
        self.trading_days_one_year = 365  # 一年的交易日数量
        self.risk_free_rate = 2.0  # 无风险利率
        self.equity_curver_df: pd.DataFrame = None  # 保存资产曲线
        self.btc_df: pd.DataFrame = None  # 保存btc的价格曲线
        self.trade_log_df: pd.DataFrame = None  # 保存交易日志
        self.performance_map = {}  # 保存回测表现
        self.year_list = []  # 回测的年份列表

    def get_R_Th_list(self):
        if self.frequency == '1d':
            self.R_list = [3, 5, 7, 10, 15, 20, 25, 30, 40, 50, 60, 75, 90, 100, 120, 150, 180, 200, 250, 300, 365, 400,
                           450, 500]
            self.Th_list = [-0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7,
                            0.8, 0.9]
        elif self.frequency == '3h':
            self.R_list = [3, 5, 10, 15, 25, 40, 60, 80, 100, 120, 160, 200, 240, 320, 400,
                           480, 640, 800, 1000, 1200, 1600, 2000, 2500, 3000]
            self.Th_list = [-0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7,
                            0.8, 0.9]
        elif self.frequency == '1h':
            self.R_list = [3, 5, 10, 15, 20, 50, 75, 120, 180, 240, 360, 480, 600, 720, 960,
                           1200, 1440, 1920, 2400, 3000, 3600, 4320, 4800, 6000]
            self.Th_list = [-0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7,
                            0.8, 0.9]

    def run(self):
        if self.frequency == "1h":
            path = "BTC_kline_1h_merge.csv"
        elif self.frequency == "3h":
            path = "BTC_kline_3h_merge.csv"
        elif self.frequency == "1d":
            path = "BTC_kline_1d_merge.csv"
        else:
            raise Exception("Frequency must be 1h or 3h or 1d!")
        self.btc_df = pd.read_csv(path)
        self.btc_df["datetime"] = pd.to_datetime(self.btc_df["datetime"])
        self.btc_df["date"] = self.btc_df["datetime"].dt.date
        self.btc_df = self.btc_df[self.btc_df['date'] >= self.data_begin_date]
        self.btc_df = self.btc_df[self.btc_df['date'] <= self.backtest_end_date]
        origin_columns = self.btc_df.columns.tolist()
        for R in self.R_list:
            for Th in self.Th_list:
                print(f"frequency:{self.frequency}, R:{R}, Th:{Th}开始回测")
                self.backtest(R, Th)
                # 回测完成后删除回测过程中添加的列
                self.btc_df = self.btc_df[origin_columns]
        self.save_performance()

    def backtest(self, R, Th):
        self.get_buy_sell_signals(R, Th)
        self.relative_amplitude_strategy()
        self.get_performance_by_begin_date(R, Th)
        self.get_performance_by_years(R, Th)
        self.save_file(R, Th)

    def get_performance_by_begin_date(self, R, Th):
        """
        按照开始日期计算表现
        :return:
        """
        for begin_date in self.begin_date_list:
            begin_datetime = datetime.datetime.strptime(begin_date, "%Y%m%d").date()
            equity_curve_df = self.equity_curver_df[self.equity_curver_df["date"] >= begin_datetime]
            hold_btc_df = self.btc_df[self.btc_df["date"] >= begin_datetime]
            trade_log_df = self.trade_log_df[self.trade_log_df["date"] >= begin_datetime]
            if trade_log_df.empty or len(trade_log_df) < 2:
                print(f"frequence:{self.frequency}, begin_date:{begin_date}没有发生交易")
                no_performance = Performance(annualized_return=np.nan,
                                             annualized_return_no_fee=np.nan,
                                             annualized_volatility=np.nan,
                                             annualized_shape=np.nan,
                                             turnover=np.nan,
                                             margin=np.nan,
                                             fitness=np.nan,
                                             excess_return=np.nan,
                                             relative_return=np.nan,
                                             information_ratio=np.nan,
                                             max_drawdown=np.nan,
                                             calmar_ratio=np.nan,
                                             yearly_count=np.nan,
                                             winning_rate=np.nan,
                                             profit_loss_ratio=np.nan,
                                             kelly_fraction=np.nan)
                self.performance_map[(begin_date, R, Th)] = no_performance
            else:
                if self.frequency in ("1h", "3h"):
                    equity_curve_df = equity_curve_df.groupby("date").last()
                    hold_btc_df = hold_btc_df.groupby("date").last()
                nav0_ser = equity_curve_df['equity0'] / equity_curve_df['equity0'].iloc[0]
                nav2_ser = equity_curve_df['equity2'] / equity_curve_df['equity2'].iloc[0]
                hold_btc_nav_ser = hold_btc_df["close"] / hold_btc_df["close"].iloc[0]
                # margin的计算需要用nav0(不含手续费) 其他指标的计算需要用nav2(双边千分之2)
                performance = self.cal_performance(nav2_ser, hold_btc_nav_ser, trade_log_df)
                performance2 = self.cal_performance(nav0_ser, hold_btc_nav_ser, trade_log_df)
                performance.margin = performance2.margin
                performance.annualized_return_no_fee = performance2.annualized_return
                self.performance_map[(begin_date, R, Th)] = performance

    def get_performance_by_years(self, R, Th):
        """
        逐年分析
        :param nav_df:
        :param trade_log_df:
        :param hold_btc_df:
        :param performance_map:
        :return:
        """
        self.equity_curver_df['year'] = self.equity_curver_df['date'].apply(lambda x: x.year)
        self.btc_df['year'] = self.btc_df['date'].apply(lambda x: x.year)
        self.trade_log_df['year'] = self.trade_log_df['date'].apply(lambda x: x.year)

        for year in self.equity_curver_df["year"].unique():
            if year not in self.year_list:
                self.year_list.append(year)
            year_equity_curve_df = self.equity_curver_df[self.equity_curver_df["year"] == year]
            year_hold_btc_df = self.btc_df[self.btc_df['year'] == year]
            if self.frequency in ("1h", "3h"):
                year_equity_curve_df = year_equity_curve_df.groupby("date").last()
                year_hold_btc_df = year_hold_btc_df.groupby("date").last()
            nav0_ser = year_equity_curve_df['equity0'] / year_equity_curve_df['equity0'].iloc[0]
            nav2_ser = year_equity_curve_df['equity2'] / year_equity_curve_df['equity2'].iloc[0]
            hold_btc_nav_ser = year_hold_btc_df['close'] / year_hold_btc_df['close'].iloc[0]
            year_trade_log_df = self.trade_log_df[self.trade_log_df['year'] == year]
            if year_trade_log_df.empty or len(year_trade_log_df) < 2:
                print(f"frequence:{self.frequency}, year:{year}没有发生交易")
                no_performance = Performance(annualized_return=np.nan,
                                             annualized_volatility=np.nan,
                                             annualized_shape=np.nan,
                                             turnover=np.nan,
                                             margin=np.nan,
                                             fitness=np.nan,
                                             excess_return=np.nan,
                                             relative_return=np.nan,
                                             information_ratio=np.nan,
                                             max_drawdown=np.nan,
                                             calmar_ratio=np.nan,
                                             yearly_count=np.nan,
                                             winning_rate=np.nan,
                                             profit_loss_ratio=np.nan,
                                             kelly_fraction=np.nan)
                self.performance_map[(year, R, Th)] = no_performance
            else:
                # 计算margin 需要用nav0的净值曲线(无手续费), 其他指标需要nav2的净值曲线(双边千2)
                performance = self.cal_performance(nav2_ser, hold_btc_nav_ser, year_trade_log_df)
                performance2 = self.cal_performance(nav0_ser, hold_btc_nav_ser, year_trade_log_df)
                performance.margin = performance2.margin
                performance.annualized_return_no_fee = performance2.annualized_return
                self.performance_map[(year, R, Th)] = performance

    def get_buy_sell_signals(self, R, Th):
        """
        得到买卖信号
        :return:
        """
        self.btc_df["high_R"] = self.btc_df["close"].rolling(R).max()
        self.btc_df["low_R"] = self.btc_df["close"].rolling(R).min()
        self.btc_df["relative_amplitude"] = (self.btc_df["close"] - self.btc_df["close"].shift(R-1)) / (
                self.btc_df["high_R"] - self.btc_df["low_R"])
        self.btc_df.replace([np.inf, -np.inf], np.nan, inplace=True)
        # 突破阈值为买入信号
        cond1 = self.btc_df["relative_amplitude"].shift(1) <= Th  # r%之前小于阈值
        cond2 = self.btc_df["relative_amplitude"] > Th  # r%现在大于阈值
        self.btc_df.loc[cond1 & cond2, ["signal"]] = 1
        # 跌破阈值为卖出信号
        cond3 = self.btc_df["relative_amplitude"].shift(1) >= Th  # r%之前大于阈值
        cond4 = self.btc_df["relative_amplitude"] < Th  # r%现在小于阈值
        self.btc_df.loc[cond3 & cond4, "signal"] = -1
        self.btc_df["signal"].fillna(0, inplace=True)

    def relative_amplitude_strategy(self):
        # 初始化仓位和净值
        position = 0  # 0表示空仓，1表示多仓
        # 初始资金10000
        equity0 = [10000]  # 没有手续费
        equity1 = [10000]  # 双边手续费千分之一
        fee_rate1 = 0.001
        equity2 = [10000]  # 双边手续费千分之二
        fee_rate2 = 0.002
        trade_log = []
        pre_price = 0
        truncate_btc_df = self.btc_df[self.btc_df['date'] >= self.backtest_begin_date]
        for _, row in truncate_btc_df.iterrows():
            # 开仓条件: 突破上轨且没有仓位
            if position == 0 and row['signal'] == 1:
                position = 1
                entry_price = row["close"]
                trade_log.append({
                    'datetime': row["datetime"],
                    'action': 'BUY',
                    'price': entry_price,
                })
                equity0.append(equity0[-1])
                equity1.append(equity1[-1] * (1 - fee_rate1))
                equity2.append(equity2[-1] * (1 - fee_rate2))
            # 平仓条件：回落到中轨且当前有多仓
            elif position == 1 and row['signal'] == -1:
                position = 0
                exit_price = row["close"]
                equity0.append(equity0[-1] * (row["close"] / pre_price))
                equity1.append(equity1[-1] * (row["close"] / pre_price) * (1 - fee_rate1))
                equity2.append(equity2[-1] * (row["close"] / pre_price) * (1 - fee_rate2))
                trade_log.append({
                    'datetime': row["datetime"],
                    'action': 'SELL',
                    'price': exit_price,
                    'return0': (exit_price - entry_price) / entry_price,
                    'return1': (exit_price - entry_price - (entry_price + exit_price) * fee_rate1) / entry_price,
                    'return2': (exit_price - entry_price - (entry_price + exit_price) * fee_rate2) / entry_price,
                })
            # 持有仓位时更新净值
            elif position == 1:
                assert pre_price > 0
                equity0.append(equity0[-1] * (row["close"] / pre_price))
                equity1.append(equity1[-1] * (row["close"] / pre_price))
                equity2.append(equity2[-1] * (row["close"] / pre_price))
            else:
                equity0.append(equity0[-1])
                equity1.append(equity1[-1])
                equity2.append(equity2[-1])
            pre_price = row["close"]
        equity0.pop(0)
        equity1.pop(0)
        equity2.pop(0)
        # 转换为DataFrame
        self.equity_curver_df = pd.DataFrame({"equity0": equity0,
                                              "equity1": equity1,
                                              "equity2": equity2})
        self.equity_curver_df["datetime"] = truncate_btc_df["datetime"].reset_index(drop=True)
        self.equity_curver_df["date"] = truncate_btc_df["date"].reset_index(drop=True)
        self.trade_log_df = pd.DataFrame(trade_log)
        if not self.trade_log_df.empty:
            self.trade_log_df["date"] = self.trade_log_df["datetime"].dt.date
        else:
            self.trade_log_df = pd.DataFrame(columns=['datetime', 'action', 'price', 'return', 'date'])

    def cal_performance(self, nav_series, nav_series_benchmark, trade_log_df):
        """
        计算策略表现
        :param nav_series:  策略每日净值曲线
        :param nav_series_benchmark:  基准每日净值曲线
        :param trade_log_df:  期间的交易记录
        :return:
        """
        performance_obj = Performance()  # 创建一个performance对象
        # 日收益率
        daily_returns = nav_series.pct_change().dropna()
        benchmark_returns = nav_series_benchmark.pct_change().dropna()
        # 总收益率
        total_return = nav_series.iloc[-1] / nav_series.iloc[0] - 1
        # 年化收益率
        n_days = len(daily_returns)  # 回测期间的交易天数
        annualized_return = (1 + total_return) ** (self.trading_days_one_year / n_days) - 1
        annualized_return = annualized_return * 100
        performance_obj.annualized_return = annualized_return
        # 年化波动率
        annualized_volatility = daily_returns.std() * np.sqrt(self.trading_days_one_year)
        annualized_volatility = annualized_volatility * 100
        performance_obj.annualized_volatility = annualized_volatility
        # 夏普比率 risk_free_rate=2.0
        annualized_shape = (annualized_return - self.risk_free_rate) / annualized_volatility
        performance_obj.annualized_shape = annualized_shape
        # 平均年交易次数
        yearly_count = len(trade_log_df) / (n_days / 365)
        performance_obj.yearly_count = yearly_count
        turnover = yearly_count / 365 * 100
        performance_obj.turnover = turnover

        # margin 以千分比的形式呈现  Margin=DailyReturn/TurnOver
        # 用年化return求daily_return
        daily_return = (annualized_return / 100 + 1) ** (1 / 365) - 1
        margin = daily_return * 100 / turnover  # 乘以100是为了和turnover保持单位统一
        # fitness, Fitness=Annualized Sharpe×Sqrt(Abs(Margin))
        fitness = annualized_shape * np.sqrt(abs(margin))
        # 先计算fitness再以千分比的形式呈现
        margin = margin * 1000
        performance_obj.margin = margin
        performance_obj.fitness = fitness

        # max_drawdown
        peak = nav_series.expanding().max()
        drawdown = (nav_series - peak) / peak
        max_drawdown = abs(drawdown.min())
        max_drawdown = max_drawdown * 100
        performance_obj.max_drawdown = max_drawdown

        # 超额收益指标
        benchmark_annualized_return = (nav_series_benchmark.iloc[-1] / nav_series_benchmark.iloc[0]) ** (
                self.trading_days_one_year / n_days) - 1
        benchmark_annualized_return = benchmark_annualized_return * 100  # benchmark的年化收益率
        performance_obj.excess_return = annualized_return - benchmark_annualized_return

        # 年化相对收益倍数
        performance_obj.relative_return = (annualized_return / 100 + 1) / (benchmark_annualized_return / 100 + 1)

        # information_ration
        # benchmark_return 是持有BTC的每日return序列 trading_days=365
        excess_returns = daily_returns.reset_index(drop=True) - benchmark_returns.reset_index(drop=True)
        excess_returns_percent = excess_returns * 100
        information_ratio = performance_obj.excess_return / (
                excess_returns_percent.std() * np.sqrt(self.trading_days_one_year))
        performance_obj.information_ratio = information_ratio

        # calmar ratio, Calmar Ratio = Annualized Return/Max Drawdown
        calmar_ratio = annualized_return / abs(max_drawdown)  # 上下都乘了100 不用处理
        performance_obj.calmar_ratio = calmar_ratio

        # 胜率
        sell_df = trade_log_df[trade_log_df['action'] == 'SELL']
        # trading_counts 一次完整的买卖为一次交易
        trading_counts = len(trade_log_df) / 2
        winning_rate = len(sell_df[sell_df["return2"] > 0]) / trading_counts
        performance_obj.winning_rate = winning_rate * 100

        # 盈亏比
        win_df = sell_df[sell_df["return2"] > 0]
        lose_df = sell_df[sell_df["return2"] < 0]
        performance_obj.profit_loss_ratio = abs(win_df["return2"].mean() / lose_df["return2"].mean())

        # 用凯利公式计算的合理配置仓位
        performance_obj.kelly_fraction = (performance_obj.profit_loss_ratio * winning_rate - (1 - winning_rate)) / (
            performance_obj.profit_loss_ratio)
        return performance_obj

    def save_file(self, R, Th):
        """
        保存所有文件
        :param R:
        :param Th:
        :return:
        """
        for begin_date in self.begin_date_list:
            begin_datetime = datetime.datetime.strptime(begin_date, "%Y%m%d").date()
            equity_curve_df = self.equity_curver_df[self.equity_curver_df["date"] >= begin_datetime]
            hold_btc_df = self.btc_df[self.btc_df["date"] >= begin_datetime]
            trade_log_df = self.trade_log_df[self.trade_log_df["date"] >= begin_datetime]
            trade_log_df.to_csv(f"{self.frequency}/{begin_date}/trade_log/trade_log_{R}_{float(Th)}.csv", index=False)
            if self.frequency in ("1h", "3h"):
                equity_curve_df = equity_curve_df.groupby("date").last()
                hold_btc_df = hold_btc_df.groupby("date").last()
                hold_btc_df.reset_index(drop=False, inplace=True)
            nav0_ser = equity_curve_df['equity0'] / equity_curve_df['equity0'].iloc[0]
            nav1_ser = equity_curve_df['equity1'] / equity_curve_df['equity1'].iloc[0]
            nav2_ser = equity_curve_df['equity2'] / equity_curve_df['equity2'].iloc[0]
            hold_btc_nav_ser = hold_btc_df["close"] / hold_btc_df["close"].iloc[0]
            nav_df = pd.DataFrame({"date": hold_btc_df["date"].values,
                                   "nav0": nav0_ser.values,
                                   "nav1": nav1_ser.values,
                                   "nav2": nav2_ser.values,
                                   "nav(hold_btc)": hold_btc_nav_ser.values})
            nav_df.to_csv(f"{self.frequency}/{begin_date}/nav/nav_{R}_{float(Th)}.csv", index=False)

    def save_performance(self):
        # 调整一下指标的顺序
        columns = ["R", "Th"] + Performance.get_all_columns()
        for begin_date in self.begin_date_list:
            results = []
            for (time, R, Th), value in self.performance_map.items():
                if time == begin_date:
                    performance_dict = value.dict()
                    performance_dict.update({'R': R, 'Th': Th})
                    results.append(performance_dict)
            results_df = pd.DataFrame(results)
            results_df = results_df[columns]
            results_df.sort_values(by=['R', 'Th'], inplace=True)
            results_df.to_csv(f"{self.frequency}/{begin_date}/result/backtest_performance.csv", index=False)

        all_df = pd.DataFrame()
        for year in self.year_list:
            results = []
            for (time, R, Th), value in self.performance_map.items():
                if time == year:
                    performance_dict = value.dict()
                    performance_dict.update({'R': R, 'Th': Th})
                    results.append(performance_dict)
            results_df = pd.DataFrame(results)
            results_df.sort_values(by=['R', 'Th'], inplace=True)
            results_df['year'] = year
            results_df = results_df[['year'] + columns]
            all_df = pd.concat([all_df, results_df])
        all_df.reset_index(drop=True, inplace=True)
        all_df.to_csv(f"{self.frequency}/backtest_performance_by_year.csv", index=False)


def delete_all_folders(current_dir="."):
    """
    删除当前目录下的所有文件夹（保留文件）
    """
    current_path = Path(current_dir)

    if not current_path.exists():
        print("当前路径不存在！")
        return

    folders_deleted = 0
    errors = 0

    # 遍历当前目录下的所有项目
    for item in current_path.iterdir():
        if item.is_dir():
            try:
                # 使用 shutil.rmtree 删除文件夹及其内容
                shutil.rmtree(item)
                print(f"✅ 已删除文件夹: {item.name}")
                folders_deleted += 1
            except Exception as e:
                print(f"❌ 删除文件夹失败 {item.name}: {e}")
                errors += 1

    print(f"\n操作完成！删除了 {folders_deleted} 个文件夹，{errors} 个错误")


def main():
    print(datetime.datetime.now())
    if not os.path.exists("BTC_kline_3h_merge.csv"):
        generate_3h_kline()
    delete_all_folders()
    for frequency in ["1d", "3h", "1h"]:
        for begin_date in ["20170101", "20200101", "20230101"]:
            os.makedirs(f"{frequency}/{begin_date}/trade_log")
            os.makedirs(f"{frequency}/{begin_date}/nav")
            os.makedirs(f"{frequency}/{begin_date}/png")
            os.makedirs(f"{frequency}/{begin_date}/result")
        backtest_obj = RelativeAmplitudeStrategy(frequency)
        backtest_obj.run()
        for begin_date in ["20170101", "20200101", "20230101"]:
            gen_html(frequency, begin_date, "R", "Th", "2025-08-31")
    print(datetime.datetime.now())


if __name__ == "__main__":
    main()
